#SLICS
#MOST UP TO DATE VERSION
#USE THIS VERSION
from threading import Thread
import math
import random
import socket
import sys
import pickle

#global "constants"
MIN_CHAR = ord(' ')
MAX_CHAR = ord('~')
PAD_LENGTH = 4096

#SEND PUBLIC KEY AS STRINGS

#needs send and receive func for keys
#send public key p
#functions for cryptography
def primeSieve(n):
    p=2
    total=[True]*n
    total[0]=False
    total[1]=False
    i=0
    while i<=n:

        if i<len(total) and total[i]==True:
            for j in range(i*i,n,i):
                total[j]=False
        i=i+1

    c=0
    prime=[]
    while c<n:
        if total[c]==True:
            #  print c,
            prime.append(c)
        c+=1
    #prime is an array of all primes less than n
    x=random.randint(5000, 50000)
    i=random.randint(5000, 50000)

    a=prime[x]
    b=prime[i]

    return ((a, b))

def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)
    "http://pastebin.com/93zxVWwS"

def gcd (a, b):

  if b == 0: return a
  else: return gcd(b, a % b)

def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m

def generateRSAKeys(p, q):

  n = p * q

  m = (p - 1) * (q - 1)

  e = random.randint(1, m-1)

  while 1:

      if gcd(m, e) == 1: break
      else: e+= 1


  d = modinv(e, m)


  return ((n,e), (n,d))

def modulo(encrypted_num, d, n):
    x=1
    y=encrypted_num
    while d>0:
        if d%2==1:
            x=(x*y)%n
        y=(y*y)%n
        d/=2
    #x=decrypted numerical value
    return x%n

def startEncrypt():
    (a, b)=primeSieve(5000000)
    (publickey, privatekey) = generateRSAKeys(a, b)
    n, e = publickey
    n, d = privatekey
    print "Public: ", publickey
    print "Private: ", privatekey
    return ((e, d, n))

def encrypt(input_num, e, c):
    encrypted_num = modulo(input_num,e,c)

    return encrypted_num

def decrypt(pickled, d, n):
    depickled = pickle.loads(pickled)#pickle
    decryptArray=[]
    decryptStr=""
    for x in depickled:
        decrypted_num=modulo(x, d, n)
        recharred=reChar(decrypted_num)

        decryptStr+=recharred

    return decryptStr

def numVal(n, eP, cP):
    length=len(n)
    enArray=[]
    for i in range(0, length):
        input_num=ord(n[i])
        encrypted_num=encrypt(input_num, eP, cP)
        enArray.append(encrypted_num)
    pickled = pickle.dumps(enArray)
    return pickled

def reChar(n):
    decrypt=chr(n)
    return decrypt

def sendKey(client_socket, e, c):
    client_socket.send(str(e))
    client_socket.send(str(c))

def recvKey(client_socket):
    e=client_socket.recv(4096)
    c=client_socket.recv(4096)
    return ((int(e), int(c)))

def authenticateClient(client_socket, cP, eP, d, c):
    data=buffRecv(client_socket)
    depickled = pickle.loads(data)
    certificate=""

    for x in depickled:
        decrypted_num=modulo(x, eP, cP)
        recharred=reChar(decrypted_num)
        certificate+=recharred
    if certificate== "Did you get this?":
        data="yes"
        reCertificate=numVal(data, d, c)
        client_socket.send(reCertificate)
    #print certificate
    else: return 0

def authenticateServer(client_socket, d, c, eP, cP):
    mssg="Did you get this?"
    certificate=numVal(mssg, d, c)
    client_socket.send(certificate)
    data=buffRecv(client_socket)#receive other certificate
    depickled = pickle.loads(data)
    reCertificate=""
    for x in depickled:
        decrypted_num=modulo(x, eP, cP)
        recharred=reChar(decrypted_num)
        reCertificate+=recharred

    if reCertificate=="yes":
        return "go"
    else: return 0

def sendPad(client_socket, cP, eP, pad):
    #enPad=numVal(pad, eP, cP)
    enPad=[0]*len(pad)
    for i in range(0, len(pad)):
        #print "pad[i]:", pad[i]
        enPad[i]= encrypt(pad[i],eP,cP)
    pickled=pickle.dumps(enPad)
    client_socket.send(pickled)
    

def receivePad(client_socket, c, d):
    enPad=buffRecv(client_socket)
    depickled=pickle.loads(enPad)    
    pad=[]
    decryptStr=""
    for x in range(0,len(depickled)):
        pad.append(modulo(depickled[x], d, c))

    return pad
    

def rotDecipher(mssg,pad):
    count =0
    Demssg=[0]*len(mssg)
    message=""
    for i in range(0, len(mssg)):
        tempC=ord(mssg[i])-pad[count]
        count+=1
        while tempC<MIN_CHAR:
            tempC=MAX_CHAR-(MIN_CHAR-tempC)#wrapping
        message+=chr(tempC)
        if count == len(pad):
            count=0
    return message
    
def rotCipher(mssg, pad):
    count =0
    Enmssg=[0]*len(mssg)
    for i in range(0, len(mssg)):
        Enmssg[i]=chr(ord(mssg[i])+pad[count])
        count+=1
        if ord(Enmssg[i])>MAX_CHAR:
            Enmssg[i]=chr(MIN_CHAR+(ord(Enmssg[i])-MAX_CHAR))#wrapping
        if count == len(pad):
            count=0
    
    return Enmssg

def genPad(length):
    pad=[0]*length
    for i in range(0, length):
        pad[i]=random.randint(MIN_CHAR, MAX_CHAR)#printable characters

    return pad

#funcs for chat system
def main():
    print "Generating public and private keys...."
    (e, d, c)= startEncrypt()
    username= raw_input("Enter your username.")
    password=raw_input("Enter your password.")
    response= password_client(e, d, c, username,password)
    print "client returned"
    if response==False:
        print "Password was incorrect, exiting program."
        return ""
    print "Password correct"
    q = raw_input("1 for client, 2 for server: ")
    if (q == '1'):
        client(e, d, c)

    elif (q == '2'):
        server(e, d, c)

    else:
        print "try again"
def password_client(e, d, c,username,password):
    print "New Client"
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    ipaddress = "localhost"#raw_input("What is the ip address of the server?")
    print "Connecting"
    client_socket.connect((ipaddress, 6677))
    print "Connected"
    (eP, cP)=recvKey(client_socket)
    sendKey(client_socket, e, c)
    authenticateClient(client_socket, cP, eP, d, c)#other user's public key
    passwordValid=[]
    t8=Thread(target=password_receive, args=(client_socket, d, c,passwordValid))
    t9=Thread(target=password_send, args=(client_socket, eP, cP,username,password,))
    print "."
    t8.start()
    t9.start()
    while 1:
        t8.join()
        t9.join()
        break
    return passwordValid[0]
def password_send(client_socket, eP, cP,username,password,):
    username=numVal(username,eP,cP)
    password=numVal(password,eP,cP)
    client_socket.send(username)
    client_socket.send(password)
    print "."

def password_receive(client_socket, d, c,passwordValid):

    while 1:
        data = client_socket.recv(92874)

        if data!="":
            print "\nReceived:", data
            if data=="True":
                passwordValid.append(True)
            else:
                passwordValid.append(False)
            break
def client(e, d, c):
    print "New Client"
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    ipaddress = "144.118.146.198" #raw_input("What is the ip address of the server?")
    print "Connecting"
    client_socket.connect((ipaddress, 3576))
    print "Connected"
    (eP, cP)=recvKey(client_socket)
    sendKey(client_socket, e, c)
    authenticateClient(client_socket, cP, eP, d, c)#other user's public key
    t3=Thread(target=receive, args=(client_socket, d, c,))
    t4=Thread(target=send, args=(client_socket, eP, cP,))
    t3.start()
    t4.start()
    while 1:
        t3.join()
        t4.join()
        break;

def server(e, d, c):
    client_socket= socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print "Binding"
    client_socket.bind(("",3576))
    print "Listening"
    client_socket.listen(5)
    comm_client, address=client_socket.accept()
    print "Connection from:", address
    sendKey(comm_client, e, c)
    (eP, cP)=recvKey(comm_client)
    go=authenticateServer(comm_client, d, c, eP, cP)#private key
    if go=="go":
        t5=Thread(target=send, args=(comm_client,eP, cP,))
        t6=Thread(target=receive, args=(comm_client, d, c,))#passing other client's public key along with your private key
        t5.start()
        t6.start()
        while 1:
            t5.join()
            t6.join()
            break;
    else: return 0

def send(client_socket, eP, cP):
    while 1:
        pad=genPad(PAD_LENGTH)
        sendPad(client_socket, cP, eP, pad) 
        data = raw_input ( "SEND (Type \quit to Quit):")
        doto=rotCipher(data, pad)
        pickled = pickle.dumps(doto)
        if (data <> '\quit' and data <> '\quit'):
            client_socket.send(pickled)
        else:
            client_socket.send(pickled)
            client_socket.close()
            break;

def receive(client_socket, d, c):
    while 1:
        pad=receivePad(client_socket, c,d)
        data=buffRecv(client_socket)
        depickled=pickle.loads(data)
        doto=rotDecipher(depickled,pad)
    
        if ( data == '\quit' or data == '\quit'):
            client_socket.close(data)
            break;
        else:

            print "\nReceived:", doto
            print "SEND (Type \quit to Quit):",

def buffRecv(client_socket):
    data =""
    index = -1
    while index < 0:
        temp = client_socket.recv(4096)
        index = temp.rfind("a.")
        if index >= 0:
            data += temp[:index+2]
        else:
            data += temp
    return data

            
main()

"""
Sources:
    http://docs.python.org
    Professor William Mongan
    Learning to Program Using Python by Cody Jackson
"""
